#!/bin/bash

set -e
pushd $(cd $(dirname ${0})/..; pwd) > /dev/null

entries=(
    "https://github.com/llvm/llvm-project.git|main|./third_party/source/mlir/llvm-project|mlir"
    "https://github.com/openxla/stablehlo.git|main|./third_party/source/mlir/stablehlo|stablehlo"
    "https://github.com/onnx/onnx-mlir.git|main|./third_party/source/mlir/onnx-mlir|"
    "https://github.com/llvm/torch-mlir.git|main|./third_party/source/mlir/torch-mlir|"
    "https://github.com/triton-lang/triton|main|./third_party/source/mlir/triton|"
    "https://github.com/tensorflow/tensorflow.git|master|./third_party/source/mlir/tensorflow|tensorflow/compiler/mlir"
    "https://github.com/tensorflow/runtime.git|master|./third_party/source/mlir/runtime|"
    "https://github.com/tensorflow/mlir-hlo.git|main|./third_party/source/mlir/mlir-hlo|"
    "https://github.com/iree-org/iree.git|main|./third_party/source/mlir/iree|compiler"
    "https://github.com/monellz/FlashTensor.git|main|./third_party/source/mlir/FlashTensor|"
    "https://github.com/plaidml/plaidml.git|plaidml-v1|./third_party/source/mlir/plaidml|pmlc"
    "https://github.com/sophgo/tpu-mlir.git|master|./third_party/source/mlir/tpu-mlir|include"
    "https://github.com/spcl/mlir-dace.git|main|./third_party/source/mlir/mlir-dace|"
    "https://github.com/woxjro/lltz.git|master|./third_party/source/mlir/lltz|"
    "https://github.com/pengmai/lagrad.git|main|./third_party/source/mlir/lagrad|include"
)

clean() {
    echo "mlir clean"
    for entry in "${entries[@]}"; do
        IFS='|' read -r url branch dir root <<< "${entry}"
        rm -rf "${dir}"
    done
}

sync() {
    echo "mlir sync"
    for entry in "${entries[@]}"; do
        IFS='|' read -r url branch dir root <<< "${entry}"
        if [ ! -d "${dir}" ]; then
            if [ -n "${root}" ]; then
                git clone --quiet --depth=1 --branch "${branch}" --single-branch --filter=blob:none --sparse "${url}" "${dir}" >/dev/null 2>&1
                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
            else
                git clone --quiet --depth=1 --branch "${branch}" --single-branch "${url}" "${dir}"
            fi
        else
            if [ -n "${root}" ]; then
                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
            fi
            git -C "${dir}" fetch --quiet --depth=1 origin "${branch}"
            git -C "${dir}" reset --quiet --hard FETCH_HEAD >/dev/null 2>&1
            git -C "${dir}" gc --quiet --prune=all
        fi
    done
}

schema() {
    echo "mlir schema"
    node ./tools/mlir_script.js
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "schema") schema;;
    esac
done
